<#
    File: Invoke-DscVmExtension.ps1
    Author: Jake Karnes (@jakekarnes42), NetSPI - 2021
    Description: PowerShell function for deploying DSC configurations hosted at any publicly accessible URL
#>

Function Invoke-DscVmExtension
{
<#
    .SYNOPSIS
        PowerShell function for executing prepackaged DSC configurations on an Azure VM. 
	.DESCRIPTION
        The function will apply a DSC configuration to an Azure VM. This is very similar to the existing Set-AzVMDscExtension. The key difference is that this function allows the DSC configuration blob to be hosted at any URL with public access. 	
#>

    [CmdletBinding()]
    Param(
        
        [Parameter(Mandatory=$true,
        ValueFromPipelineByPropertyName = $true,
        HelpMessage="The name of the VM to deploy the DSC to.")]
        [Alias('Name')]
        [string]  $VMName,
        
        [Parameter(Mandatory=$true,
        ValueFromPipelineByPropertyName = $true,
        HelpMessage="The name of the resource group the VM belongs to.")]
        [string]  $ResourceGroupName,
     
        [Parameter(Mandatory=$true,
        HelpMessage="The publicly accessible URL hosting the DSC archive. This may be a storage account URL generated by Publish-AzVMDscConfiguration cmdlet after publishing, if the blob allows public access. To create archive to host elsewhere, use the Publish-AzVMDscConfiguration cmdlet with the -OutputArchivePath option to generate the archive locally. It is expected that the URL will end with the archive file's name (e.g. 'https://[MY-STORAGE-ACCOUNT].blob.core.windows.net/windows-powershell-dsc/ExampleDSCArchive.ps1.zip')")]
        [string]  $ConfigurationArchiveURL,
 
        [Parameter(Mandatory=$false,
        HelpMessage="The DSC configuration to be deployed. This is the name of the DSC configuration within the configuration archive.")]
        [string] $ConfigurationName = [System.IO.Path]::GetFileNameWithoutExtension([System.IO.Path]::GetFileNameWithoutExtension((([uri] $ConfigurationArchiveURL).Segments[-1]))), #Assume the Configuration has the same name as the ZIP file (excluding the .ps1 and .zip extensions)

        [Parameter(Mandatory=$false,
        HelpMessage="The configuration arguments to be passed to the DSC configuration, if any.")]
        [hashtable] $ConfigurationArgument = @{},
        
        [Parameter(Mandatory=$false,
        HelpMessage="The authorization context with permissions to deploy VM extensions. By default, uses the current context from `Get-AzContext`.")]
        [Microsoft.Azure.Commands.Profile.Models.Core.PSAzureContext] $Context
    )

    #The logic below is a reimplementation of the Set-AzVMDscExtension cmdlet
    #This is required to use an arbitrary URL, rather than a storage account. 
    #From: https://github.com/Azure/azure-powershell/blob/master/src/Compute/Compute/Extension/DSC/SetAzureVMDscExtensionCommand.cs

    #Identify the archive file name from the URL
    $ConfigurationArchive  = ([uri] $ConfigurationArchiveURL).Segments[-1]

    #Parse input configuration args
    $parsedConfigurationArguments = [Microsoft.WindowsAzure.Commands.Common.Extensions.DSC.DscExtensionSettingsSerializer]::SeparatePrivateItems($ConfigurationArgument)

    #Build the public settings object
    $publicSettings = New-Object -Type Microsoft.WindowsAzure.Commands.Common.Extensions.DSC.DscExtensionPublicSettings
    $publicSettings.ModulesUrl = $ConfigurationArchiveURL
    $publicSettings.Privacy = @{ DataCollection = "Disable" }
    $publicSettings.ConfigurationFunction = "{0}\{1}" -f [System.IO.Path]::GetFileNameWithoutExtension($ConfigurationArchive),$ConfigurationName
    $publicSettings.Properties = $parsedConfigurationArguments.Item1;

    #Build the private settings object
    $privateSettings = New-Object -Type Microsoft.WindowsAzure.Commands.Common.Extensions.DSC.DscExtensionPrivateSettings
    $privateSettings.Items = $parsedConfigurationArguments.Item2;
    $privateSettings.DataBlobUri = $null; #"ConfigurationData" not supported since it requires an upload to a storage account

    #Get the VM's location
    $Location = (Get-AzVM -Name $VMName -ResourceGroupName $ResourceGroupName).Location

    #Build the VirtualMachineExtension object
    $parameters = New-Object -Type Microsoft.Azure.Management.Compute.Models.VirtualMachineExtension -Property @{
                    Location = $Location;
                    Publisher = "Microsoft.Powershell";
                    VirtualMachineExtensionType = "DSC";
                    TypeHandlerVersion = "2.83"; #latest version as of June 5, 2021
                    Settings = $publicSettings;
                    ProtectedSettings = $privateSettings;
                    AutoUpgradeMinorVersion = $false
                };

    #Create a client using the authorized context
    if ($Context -eq $null){
        $Context = Get-AzContext
    }
    $computeClient = [Microsoft.Azure.Commands.Compute.ComputeClient]::new($Context)
    $vmExtensionClient = $computeClient.ComputeManagementClient.VirtualMachineExtensions

    #Send it
    Write-Host "Deploying DSC to VM: $VMName"
    $op = $vmExtensionClient.CreateOrUpdateWithHttpMessagesAsync($ResourceGroupName, $VMName, "Microsoft.Powershell.DSC" , $parameters).GetAwaiter().GetResult();

    # If the deployment fails, we simply don't get a response, but the above prints an error
    if($op.Response) {
        Write-Host "Deployment Successful: $($op.Response.IsSuccessStatusCode)" 
        Write-Verbose $op.Response.Content.ReadAsStringAsync().Result | ConvertFrom-Json
    } 

    
    #Automatic cleanup
    Write-Host "Deleting DSC extension from VM: $VMName"
    $op = $vmExtensionClient.DeleteWithHttpMessagesAsync($ResourceGroupName, $VMName, "Microsoft.Powershell.DSC").GetAwaiter().GetResult();

    if($op.Response) {
        Write-Host "Removal Successful: $($op.Response.IsSuccessStatusCode)" 
        Write-Verbose $op.Response.Content.ReadAsStringAsync().Result | ConvertFrom-Json
    } 

}